#!/bin/bash

# Copyright 2019 Mobvoi Inc. All Rights Reserved.
. ./path.sh || exit 1;

# Use this to control how many gpu you use, It's 1-gpu training if you specify
# just 1gpu, otherwise it's is multiple gpu training based on DDP in pytorch
export CUDA_VISIBLE_DEVICES="0,1,3,4,5,6"

stage=4 # start from 0 if you need to start from data preparation
stop_stage=4

dict=data_list/units_en_cn.txt
bpe_model=data_list/en_cn_bpe.model
data_type=shard


train_config=conf/train_whisper_medium_streaming.yaml
dir=/home/work_nfs6/xlgeng/new_workspace/wenet_gxl_en_cn/streaming_fbank_exp
tensorboard_dir=$dir/tensorboard
checkpoint=/home/work_nfs6/xlgeng/new_workspace/wenet_gxl_en_cn/streaming_fbank_exp/epoch_7.pt


HOST_NODE_ADDR="localhost:0"
num_nodes=1
job_id=202
cmvn=false
nj=16
num_workers=8
prefetch=500

# use average_checkpoint will get better result
average_checkpoint=false
decode_checkpoint=$dir/9.pt
average_num=30
decode_modes="ctc_greedy_search ctc_prefix_beam_search attention attention_rescoring"

train_engine=torch_ddp

deepspeed_config=conf/ds_stage2.json
deepspeed_save_states="model_only"

. tools/parse_options.sh || exit 1;





if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
  mkdir -p $dir
  num_gpus=$(echo $CUDA_VISIBLE_DEVICES | awk -F "," '{print NF}')
  dist_backend="nccl"
  cmvn_opts=
  $cmvn && cp data/${train_set}/global_cmvn4paraformer $dir/global_cmvn
  $cmvn && cmvn_opts="--cmvn ${dir}/global_cmvn"

  if [ ${train_engine} == "deepspeed" ]; then
    echo "$0: using deepspeed"
  else
    echo "$0: using torch ddp"
  fi

  echo "$0: num_nodes is $num_nodes, proc_per_node is $num_gpus"
  torchrun --nnodes=$num_nodes --nproc_per_node=$num_gpus \
           --rdzv_id=$job_id --rdzv_backend="c10d" --rdzv_endpoint=$HOST_NODE_ADDR \
    wenet/bin/train.py \
      --train_engine ${train_engine} \
      --config $train_config \
      --data_type  $data_type \
      --train_data data_list/asru_train.shards \
      --cv_data data_list/asru_dev.shards \
      ${checkpoint:+--checkpoint $checkpoint} \
      --model_dir $dir \
      --tensorboard_dir ${tensorboard_dir} \
      --ddp.dist_backend $dist_backend \
      --num_workers ${num_workers} \
      --prefetch ${prefetch} \
      $cmvn_opts \
      --pin_memory \
      --deepspeed_config ${deepspeed_config} \
      --deepspeed.save_states ${deepspeed_save_states}
fi

